#include "StableINSSolver2D.hh"
#include "../algorithms/Bresenham.hh"


struct AdvectionTracker {

  const bool *boundary;

  int gridSizeX, gridSizeY;

  int iEnd, jEnd;

  AdvectionTracker(const bool *boundary, int gridSizeX, int gridSizeY) {
    this->boundary = boundary;
    this->gridSizeX = gridSizeX;
    this->gridSizeY = gridSizeY;
    this->iEnd = -1;
    this->jEnd = -1;
  }

  inline bool operator()(int i, int j) {
    if (boundary[j * gridSizeX + i]) {
      return false;
    }

    iEnd = i;
    jEnd = j;
    return true;
  }
};


StableINSSolver2D::StableINSSolver2D()
: velocityX(NULL), velocityY(NULL), boundary(NULL), tmpGrid(NULL) {
}

StableINSSolver2D::~StableINSSolver2D() {
  freeData();
}

const float *StableINSSolver2D::getVelocityX() {
  return velocityX;
}

const float *StableINSSolver2D::getVelocityY() {
  return velocityY;
}

void StableINSSolver2D::resetSystem(
    int domainSizeX, int domainSizeY, float gridSpacing,
    const float *velocityX, const float *velocityY,
    const bool *boundary) {
  freeData();

  this->domainSizeX = domainSizeX;
  this->domainSizeY = domainSizeY;
  this->hx = gridSpacing;
  this->hy = gridSpacing;
  this->gridSizeX = 1 + domainSizeX + 1;
  this->gridSizeY = 1 + domainSizeY + 1;
  this->gridSize = gridSizeX * gridSizeY;

  this->velocityX  = new float[gridSize];
  this->velocityY  = new float[gridSize];
  this->tmpVelocityX  = new float[gridSize];
  this->tmpVelocityY  = new float[gridSize];
  this->boundary = new bool[gridSize];
  if (velocityX) {
    std::copy(velocityX, velocityX + gridSize, this->velocityX);
  } else {
    std::fill(this->velocityX, this->velocityX + gridSize, 0.0f);
  }
  if (velocityY) {
    std::copy(velocityY, velocityY + gridSize, this->velocityY);
  } else {
    std::fill(this->velocityY, this->velocityY + gridSize, 0.0f);
  }
  std::copy(boundary, boundary + gridSize, this->boundary);

  /* Make sure the grid boundary is marked as such. */
  for (int i = 0; i < gridSizeX; i++) {
    this->boundary[idx(i, 0)] = true;
    this->boundary[idx(i, gridSizeY - 1)] = true;
  }
  for (int j = 0; j < gridSizeY; j++) {
    this->boundary[idx(0, j)] = true;
    this->boundary[idx(gridSizeX - 1, j)] = true;
  }
}

void StableINSSolver2D::updateVelocity(
    float *sourceX, float *sourceY,
    float viscosity, float vorticity, float dt,
    BoundaryCondition bc) {
  if (bc == BC_DIRICHLET) {
    copyBoundaryValues(velocityX, sourceX);
    copyBoundaryValues(velocityY, sourceY);

    copyBoundaryValues(tmpVelocityX, sourceX);
    copyBoundaryValues(tmpVelocityY, sourceY);
  }

  advect(tmpVelocityX, velocityX, dt, bc);
  advect(tmpVelocityY, velocityY, dt, bc);

  addSource(tmpVelocityX, tmpVelocityX, sourceX, dt);
  addSource(tmpVelocityY, tmpVelocityY, sourceY, dt);

  addVorticity(tmpVelocityX, tmpVelocityY, vorticity, dt, bc);

  diffuse(velocityX, tmpVelocityX, viscosity, dt, bc);
  diffuse(velocityY, tmpVelocityY, viscosity, dt, bc);

  float *pressure = tmpVelocityX;
  float *divergence = tmpVelocityY;
  projectVelocity(pressure, divergence, dt, bc);
}

void StableINSSolver2D::transportScalarField(
    float *field, float *source,
    float diffusion, float dissipation, float dt,
    BoundaryCondition bc) {
  if (bc == BC_DIRICHLET) {
    copyBoundaryValues(field, source);
  }

  addSource(field, field, source, dt);
  advect(source, field, dt, bc);
  diffuse(field, source, diffusion, dt, bc);
  dissipate(field, dissipation, dt);
}

void StableINSSolver2D::copyBoundaryValues(float *target, const float *source) {
  #pragma omp parallel for
  for (int i = 0; i < gridSize; i++) {
    if (boundary[i]) {
      target[i] = source[i];
    }
  }
}

void StableINSSolver2D::addSource(
    float *target,
    const float *field, const float *source,
    float dt) {
  #pragma omp parallel for
  for (int i = 0; i < gridSize; i++) {
    if (!boundary[i]) {
      target[i] = field[i] + dt * source[i];
    }
  }
}

void StableINSSolver2D::addVorticity(
    float *u, float *v,
    float vorticity, float dt,
    BoundaryCondition bc) {
  if (vorticity == 0.0f) {
    return;
  }

  if (!tmpGrid) {
    tmpGrid = new float[gridSize];
  }

  float *curl = tmpGrid;

  #pragma omp parallel for
  for (int j = 1; j <= domainSizeY; j++) {
    for (int i = 1; i <= domainSizeX; i++) {
      if (!boundary[idx(i, j)]) {
        float d1 =   value(v, i, j, i + 1,   j  , bc)
                   - value(v, i, j, i - 1,   j  , bc);
        float d2 =   value(u, i, j,   i  , j + 1, bc)
                   - value(u, i, j,   i  , j - 1, bc);
        curl[idx(i, j)] = d1 / (2.0f * hx) - d2 / (2.0f * hy);
      }
    }
  }

  #pragma omp parallel for
  for (int j = 1; j <= domainSizeY; j++) {
    for (int i = 1; i <= domainSizeX; i++) {
      if (!boundary[idx(i, j)]) {
        float nx =   std::abs(value(curl, i, j, i + 1, j, BC_DIRICHLET_ZERO))
                   - std::abs(value(curl, i, j, i - 1, j, BC_DIRICHLET_ZERO));
        nx /= 2.0f * hx;

        float ny =   std::abs(value(curl, i, j, i, j + 1, BC_DIRICHLET_ZERO))
                   - std::abs(value(curl, i, j, i, j - 1, BC_DIRICHLET_ZERO));
        ny /= 2.0f * hy;

        float nl = std::sqrt(nx * nx + ny * ny);
        if (nl > 0.00001f) {
          nx /= nl;
          ny /= nl;

          float fx = vorticity * hx *  (ny * curl[idx(i, j)]);
          float fy = vorticity * hy * -(nx * curl[idx(i, j)]);

          u[idx(i, j)] += dt * fx;
          v[idx(i, j)] += dt * fy;
        }
      }
    }
  }
}

void StableINSSolver2D::diffuse(
    float *field, const float *field0,
    float factor, float dt,
    BoundaryCondition bc) {
  if (factor == 0.0f) {
    std::copy(field0, field0 + gridSize, field);
    return;
  }

  float ai = -dt * factor / (hx * hx);
  float aj = -dt * factor / (hy * hy);
  float a  = 1.0f - 2.0f * (ai + aj);
  solveLinearSystem(field, field0, a, ai, aj, bc);
}

void StableINSSolver2D::advect(
    float *field, const float *field0,
    float dt,
    BoundaryCondition bc) {
  #pragma omp parallel for
  for (int j = 1; j <= domainSizeY; j++) {
    for (int i = 1; i <= domainSizeX; i++) {
      if (!boundary[idx(i, j)]) {
        float x = i - velocityX[idx(i, j)] * dt / hx;
        float y = j - velocityY[idx(i, j)] * dt / hy;

        float s1 = std::abs(x - (int) x);
        float s0 = 1.0f - s1;
        float t1 = std::abs(y - (int) y);
        float t0 = 1.0f - t1;

        AdvectionTracker tracker(boundary, gridSizeX, gridSizeY);
        Bresenham::line2D(i, j, (int) x, (int) y, tracker);
        int i0 = tracker.iEnd;
        int j0 = tracker.jEnd;

        float v1 = field0[idx(i0, j0)];
        float v2 = value(field0, i0, j0,   i0  , j0 + 1, bc);
        float v3 = value(field0, i0, j0, i0 + 1,   j0  , bc);
        float v4 = value(field0, i0, j0, i0 + 1, j0 + 1, bc);

        field[idx(i, j)] =   s0 * (t0 * v1 + t1 * v2)
                           + s1 * (t0 * v3 + t1 * v4);
      }
    }
  }
}

void StableINSSolver2D::projectVelocity(
    float *pressure, float *divergence,
    float dt,
    BoundaryCondition bc) {
  #pragma omp parallel for
  for (int j = 1; j <= domainSizeY; j++) {
    for (int i = 1; i <= domainSizeX; i++) {
      if (!boundary[idx(i, j)]) {
        float du =   value(velocityX, i, j, i + 1,   j  , bc)
                   - value(velocityX, i, j, i - 1,   j  , bc);
        float dv =   value(velocityY, i, j,   i  , j + 1, bc)
                   - value(velocityY, i, j,   i  , j - 1, bc);

        divergence[idx(i, j)] = -(du / (2.0f * hx) + dv / (2.0f * hy));

        pressure[idx(i, j)] = 0.0f;
      }
    }
  }

  float ai = dt / (hx * hx);
  float aj = dt / (hy * hy);
  float a  = -2.0f * (ai + aj);
  solveLinearSystem(pressure, divergence, -a, -ai, -aj, BC_NEUMANN_ZERO);

  #pragma omp parallel for
  for (int j = 1; j <= domainSizeY; j++) {
    for (int i = 1; i <= domainSizeX; i++) {
      if (!boundary[idx(i, j)]) {
        float dp;

        dp =   value(pressure, i, j, i + 1, j, BC_NEUMANN_ZERO)
             - value(pressure, i, j, i - 1, j, BC_NEUMANN_ZERO);
        velocityX[idx(i, j)] -= dt * (dp / (2.0f * hx));

        dp =   value(pressure, i, j, i, j + 1, BC_NEUMANN_ZERO)
             - value(pressure, i, j, i, j - 1, BC_NEUMANN_ZERO);
        velocityY[idx(i, j)] -= dt * (dp / (2.0f * hy));
      }
    }
  }
}

void StableINSSolver2D::dissipate(float *field, float dissipation, float dt) {
  if (dissipation == 0.0f) {
    return;
  }

  #pragma omp parallel for
  for (int i = 0; i < gridSize; i++) {
    if (!boundary[i]) {
      field[i] /= 1.0f + dt * dissipation;
    }
  }
}

void StableINSSolver2D::solveLinearSystem(
    float *x, const float *x0,
    float a, float ai, float aj,
    BoundaryCondition bc) {
#if defined(JACOBI_SOLVER)
  if (!tmpGrid) {
    tmpGrid = new float[gridSize];
  }

  float *xNew = tmpGrid;

  if (bc == BC_DIRICHLET) {
    copyBoundaryValues(xNew, x);
  }

  for (int k = 0; k < 20; k++) {
    #pragma omp parallel for
    for (int j = 1; j <= domainSizeY; j++) {
      for (int i = 1; i <= domainSizeX; i++) {
        if (!boundary[idx(i, j)]) {
          float sum  =   ai * value(x, i, j, i - 1,   j  , bc)
                       + ai * value(x, i, j, i + 1,   j  , bc)
                       + aj * value(x, i, j,   i  , j - 1, bc)
                       + aj * value(x, i, j,   i  , j + 1, bc);

          xNew[idx(i, j)] = (x0[idx(i, j)] - sum) / a;
        }
      }
    }
    std::swap(x, xNew);
  }
#else
  for (int k = 0; k < 20; k++) {
    for (int j = 1; j <= domainSizeY; j++) {
      for (int i = 1; i <= domainSizeX; i++) {
        if (!boundary[idx(i, j)]) {
          float sum  =   ai * value(x, i, j, i - 1,   j  , bc)
                       + ai * value(x, i, j, i + 1,   j  , bc)
                       + aj * value(x, i, j,   i  , j - 1, bc)
                       + aj * value(x, i, j,   i  , j + 1, bc);

          x[idx(i, j)] = (x0[idx(i, j)] - sum) / a;
        }
      }
    }
  }
#endif
}

inline int StableINSSolver2D::idx(int i, int j) {
  return j * gridSizeX + i;
}

inline float StableINSSolver2D::value(
    const float *x,
    int i0, int j0,
    int i1, int j1,
    BoundaryCondition bc) {
  if (boundary[idx(i1, j1)] && (bc != BC_DIRICHLET)) {
    if (bc == BC_NEUMANN_ZERO) {
      return x[idx(i0, j0)];
    } else /* BC_DIRICHLET_ZERO */ {
      return 0.0f;
    }
  } else {
    return x[idx(i1, j1)];
  }
}

void StableINSSolver2D::freeData() {
  delete[] velocityX;
  delete[] velocityY;
  delete[] boundary;
  delete[] tmpGrid;
}

